package com.qiuzhao.handler;

import com.alibaba.nls.client.AccessToken;
import com.alibaba.nls.client.protocol.InputFormatEnum;
import com.alibaba.nls.client.protocol.NlsClient;
import com.alibaba.nls.client.protocol.SampleRateEnum;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriber;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriberListener;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriberResponse;
import com.alibaba.nls.client.protocol.tts.SpeechSynthesizer;
import com.qiuzhao.entity.Role;
import com.qiuzhao.mapper.RoleMapper;
import com.qiuzhao.service.impl.LLMService;
import com.qiuzhao.utils.AliyunTtsUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Schedulers;

import java.io.ByteArrayOutputStream;
import java.time.Duration;
import java.util.Base64;

@Component
@Slf4j
public class AliyunAsrWebSocketHandler implements WebSocketHandler {

    private static final String APP_KEY = "lityBgzjbG7JIWaA";
    private static final String ACCESS_KEY_ID = "LTAI5t73zgPFVdh963i12ehd";
    private static final String ACCESS_KEY_SECRET = "Y86n52JWBIQcQPBF3u8fznc5jLFKHy";
    private static final String ASR_URL = "wss://nls-gateway-cn-shanghai.aliyuncs.com/ws/v1";

    @Autowired
    private ChatClient chatClient;
    @Autowired
    private RoleMapper roleMapper;

    @Override
    public Mono<Void> handle(WebSocketSession session) throws RuntimeException {
        return Mono.defer(() -> {
            try {
                String query = session.getHandshakeInfo().getUri().getQuery();
                String[] split = query.split("&");
                String conversationId = split[0].split("=")[1];
                String roleId = split[1].split("=")[1];

                // 获取临时 token
                AccessToken accessToken = new AccessToken(ACCESS_KEY_ID, ACCESS_KEY_SECRET);
                accessToken.apply();

                NlsClient client = new NlsClient(ASR_URL, accessToken.getToken());
                Sinks.Many<String> resultSink = Sinks.many().multicast().onBackpressureBuffer();

                SpeechTranscriber transcriber = new SpeechTranscriber(client, getListener(resultSink, conversationId, roleId));

                transcriber.setAppKey(APP_KEY);
                transcriber.setFormat(InputFormatEnum.PCM);
                transcriber.setSampleRate(SampleRateEnum.SAMPLE_RATE_16K);
                transcriber.setEnableIntermediateResult(true);
                transcriber.setEnablePunctuation(true);
                transcriber.setEnableITN(true);

                transcriber.start();

                // 接收前端音频流并发送给阿里云
                session.receive()
                        .map(WebSocketMessage::getPayload)
                        .doOnNext(dataBuffer -> {
                            byte[] bytes = new byte[dataBuffer.readableByteCount()];
                            dataBuffer.read(bytes);
                            //清除dataBuffer
                            DataBufferUtils.release(dataBuffer);
                            transcriber.send(bytes, bytes.length);
                        })
                        .doOnComplete(() -> {
                            try {
                                transcriber.stop();
                            } catch (Exception e) {
                                log.error("[ASR Stop Error]", e);
                            }
                        })

                        .subscribe();

                // 将识别结果推送到前端
                return session.send(
                        resultSink.asFlux().map(session::textMessage)
                ).doFinally(sig -> {
                    try {
                        transcriber.close();
                        client.shutdown();
                    } catch (Exception ignore) {
                    }
                });

            } catch (Exception e) {
                return session.close();
            }
        });
    }

    private SpeechTranscriberListener getListener(Sinks.Many<String> sink, String conversationId, String roleId) {
        return new SpeechTranscriberListener() {
            @Override
            public void onTranscriberStart(SpeechTranscriberResponse response) {
                log.info("[ASR Start] taskId={}", response.getTaskId());
            }

            @Override
            public void onTranscriptionResultChange(SpeechTranscriberResponse response) {
                String result = response.getTransSentenceText();
                if (result != null && !result.isEmpty()) {
                    log.info("[中间结果] " + result);
                }
            }

            @Override
            public void onSentenceBegin(SpeechTranscriberResponse response) {
                log.info("[Sentence Begin]");
            }

            @Override
            public void onSentenceEnd(SpeechTranscriberResponse response) {
                String result = response.getTransSentenceText();
                if (result != null && !result.isEmpty()) {
                    log.info("[最终结果] " + result);

                    // 异步执行 LLM + TTS
                    Mono.fromCallable(() -> {
                                Role role = roleMapper.selectById(roleId);
                                if (role == null)
                                    throw new RuntimeException("角色不存在");
                                String content = chatClient
                                        .prompt()
                                        .system(role.getSystemPrompt())
                                        .user(result)
                                        .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, conversationId + "-" + roleId))
                                        .call()
                                        .content();

//                                sink.tryEmitNext("[大模型结果] " + content);

                                AliyunTtsUtil tts = new AliyunTtsUtil(APP_KEY, ACCESS_KEY_ID, ACCESS_KEY_SECRET, ASR_URL, role);
                                byte[] audioData = tts.synthesize(content);
                                String base64 = Base64.getEncoder().encodeToString(audioData);
                                sink.tryEmitNext("[TTS结果] " + base64);

                                return true;
                            })
                            .subscribeOn(Schedulers.boundedElastic()) // 放到独立线程池
                            .subscribe();
                }
            }


            @Override
            public void onTranscriptionComplete(SpeechTranscriberResponse response) {
                log.info("[ASR Complete]");
                sink.tryEmitComplete();
            }

            @Override
            public void onFail(SpeechTranscriberResponse response) {
                sink.tryEmitNext("[识别失败状态码] " + response.getStatus());
                sink.tryEmitNext("[识别失败状态文本] " + response.getStatusText());
                sink.tryEmitComplete();
            }
        };
    }
}
